
from patient_information import find_patient_files,get_murmur,get_grade,load_patient_data
import os
import librosa
import soundfile
import shutil
from tqdm import tqdm

# def data_cut(
#     data_directory: str,#数据集路径
#     out_directory: str,
# ):
data_directory ="/home/dsp610/HZH/2022_challenge_new/the-circor-digiscope-phonocardiogram-dataset-1.0.3/test_data"
out_directory = 'data/no_fold/vali_data'
if not os.path.exists(out_directory):
    os.makedirs(out_directory)
else:
    # exit()
    shutil.rmtree(out_directory)  # 删除输出路径下的文件夹
    os.makedirs(out_directory)
patient_files = find_patient_files(data_directory)  # 获取排序后的病人text文件
num_patient_files = len(patient_files)
for f in sorted(os.listdir(data_directory)):
    root, extension = os.path.splitext(f)
    if extension == '.txt':
        with open(os.path.join(data_directory, f)) as txt_f:
            txt_data = txt_f.read()
        murmur_unkonwn = get_murmur(txt_data)
        if murmur_unkonwn != 'Unknown':
            _ = shutil.copy(os.path.join(data_directory, f), out_directory)
    elif extension == '.wav':
        location = root.split("_")[1].strip()
        patient_ID = root.split("_")[0].strip()
        with open(os.path.join(data_directory, patient_ID + '.txt'), 'r') as txt_f:
            txt_data = txt_f.read()
        murmur_unkonwn = get_murmur(txt_data)
        if murmur_unkonwn != 'Unknown':
            grade = get_grade(txt_data)
            recording, fs = librosa.load(os.path.join(data_directory, f), sr=4000)  # 分割（3s不重叠）
            num_cut = len(recording) / (3 * 4000)  # 每个记录的片段数量
            if num_cut >= 2:
                recording = recording[2 * fs:len(recording) - fs]
            # recording = (recording- np.mean(recording))/ np.max(np.abs(recording)) #幅值归一化
            # recording = schmidt_spike_removal(recording) #去尖峰
            start = 0
            end = start + 3 * fs
            cut = list()
            num_cut = len(recording) / (3 * 4000)
            for num in range(int(num_cut)):  # 将每个片段写入对应的听诊区文件夹
                small = recording[start:end]
                cut.append(small)
                soundfile.write(
                    out_directory + '/' + patient_ID + '_' + str(location) + '_' + str(grade) + '_' + str(num) + '.wav',
                    cut[num], fs)
                start += 3 * fs
                end = start + 3 * fs

# patient_files = find_patient_files(out_directory)#获取排序后的病人text文件
# num_patient_files = len(patient_files)
# murmur_classes = ["Absent", "Soft", "Loud"]
# num_murmur_classes = len(murmur_classes)
# murmurs = list()
# patient_ID = list()
# for i in tqdm(range(num_patient_files)): #可视化遍历进度条，遍历每个病人的txt文件
#     # Load the current patient data and recordings.
#     current_patient_data = load_patient_data(patient_files[i]) #获取text文本内容
#     murmur_unkonwn=get_murmur(current_patient_data)
#     #跳过unknown的数据
#     if murmur_unkonwn == 'Unknown':
#         print(patient_files[i])
